package com.konghao.test.junit.service;

import com.konghao.test.junit.util.DbUtil;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.ext.mysql.MySqlConnection;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

import java.io.*;
import java.sql.SQLException;

/**
 * AbstractDbUnitTestCase--JunitTest
 *
 * @author: Jovi
 * @createTime: 2017-10-10 17:33
 **/
public abstract class AbstractDbUnitTestCase {

    public static IDatabaseConnection dbunitConn;

    private File tempFile;

    @BeforeClass
    public static void init() throws SQLException, DatabaseUnitException {
        dbunitConn = new MySqlConnection(DbUtil.getConnection(), "test_junit");
    }

    protected IDataSet createDateSet(String name) throws DataSetException {

        InputStream is = getClass().getResourceAsStream("/t_" + name + ".xml");

        Assert.assertNotNull("Dbunit的基本数据文件不存在", is);

        return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is)));

    }

    protected void backUpAllTable() throws SQLException, IOException, DataSetException {
        IDataSet ds = dbunitConn.createDataSet();
        writeBackupFile(ds);
    }

    private void writeBackupFile(IDataSet ds) throws IOException, DataSetException {
        tempFile = new File("src/test/resource/backup.xml");
        FlatXmlDataSet.write(ds, new FileWriter(tempFile));
    }

    protected void backupCustomTable(String[] tname) throws DataSetException, IOException {
        QueryDataSet ds = new QueryDataSet(dbunitConn);
        for (String str : tname) {
            ds.addTable(str);
        }
        writeBackupFile(ds);
    }

    protected void backupOneTable(String name) throws IOException, DataSetException {
        backupCustomTable(new String[]{name});
    }

    protected void resumeTable() throws DatabaseUnitException, SQLException, FileNotFoundException {
        IDataSet ds = new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tempFile))));

        DatabaseOperation.CLEAN_INSERT.execute(dbunitConn, ds);
    }

    @AfterClass
    public static void destory() {

        if (dbunitConn != null) {
            try {
                dbunitConn.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

}
